# -*- coding: utf-8 -*-
import uuid
from typing import List

from sqlalchemy import desc, func

from db import UserGroup
from db.session import with_session
from db.models.conversation import Conversation, Dialog


""" ============================== Conversation =============================="""


@with_session
def db_get_superman_conversation(session, dialog_id) -> List[Conversation]:
    return session.query(Conversation) \
        .join(Dialog, Dialog.id == Conversation.dialog_id) \
        .filter(Dialog.id == dialog_id, Dialog.is_delete.is_(False), Conversation.is_delete.is_(False))\
        .all()


@with_session
def db_get_admin_conversation(session, user_id, dialog_id) -> List[Conversation]:
    # 获取当前用户所在的所有组ID
    subquery_groups = session.query(UserGroup.group).filter(UserGroup.user == user_id)

    # 获取这些组中的所有用户ID
    subquery_users = session.query(UserGroup.user).filter(UserGroup.group.in_(subquery_groups))

    # 查询符合条件的对话记录
    return session.query(Conversation) \
        .join(Dialog, Dialog.id == Conversation.dialog_id) \
        .filter(Dialog.id == dialog_id, Dialog.user_id.in_(subquery_users), Dialog.is_delete.is_(False), Conversation.is_delete.is_(False)).all()


@with_session
def db_get_other_conversation(session, user_id, dialog_id) -> List[Conversation]:
    return session.query(Conversation)\
        .join(Dialog, Dialog.id == Conversation.dialog_id)\
        .filter(Dialog.id == dialog_id, Dialog.user_id == user_id, Dialog.is_delete.is_(False), Conversation.is_delete.is_(False))\
        .all()


@with_session
def db_get_conversations(session, dialog_id, user_id) -> List[Conversation]:
    conversation = session.query(Conversation) \
                            .join(Dialog, Dialog.id == Conversation.dialog_id) \
                            .filter(
                                Dialog.id == dialog_id,
                                Dialog.user_id == user_id,
                                Dialog.is_delete == False,
                                Conversation.is_delete == False
                                ) \
                            .order_by(desc(Conversation.updated_time)) \
                            .limit(3)
    return conversation


@with_session
def db_add_conversation(session, question, dialog_id, references, answer="", image_path=None, file_id=None) -> Conversation:
    # 确保references是列表类型
    if references is None:
        references = []
    elif not isinstance(references, list):
        references = [references]
    
    # 确保file_id是字符串
    if file_id is not None:
        file_id = str(file_id)
    
    conversion = Conversation(id=uuid.uuid1(),
                              question=question,
                              answer=answer,
                              dialog_id=dialog_id,
                              reference=references,
                              image_path=image_path,
                              file_id=file_id
                              )
    session.add(conversion)
    session.commit()
    return conversion


@with_session
def db_update_conversion_answer(session, conversation_id, answer: str) -> Conversation:
    conversation = session.query(Conversation).get(conversation_id)
    conversation.answer = answer
    session.commit()
    return conversation


@with_session
def db_delete_conversation(session, user_id, conversation_id):
    conversion = session.query(Conversation)\
        .join(Dialog, Dialog.id == Conversation.dialog_id)\
        .filter(Conversation.id == conversation_id, Dialog.user_id == user_id, Dialog.is_delete.is_(False), Conversation.is_delete.is_(False))\
        .first()
    if conversion is None:
        return None

    conversion.is_delete = True
    session.add(conversion)
    session.commit()
    return True


""" ============================= Dialog ============================="""


@with_session
def db_search_dialog(session, dialog_id, user_id, dialog_type=None) -> Dialog:
    # 确保dialog_id是UUID对象
    if isinstance(dialog_id, str):
        dialog_id = uuid.UUID(dialog_id)
    
    query = session.query(Dialog).filter(
        Dialog.id == dialog_id,
        Dialog.user_id == user_id,
        Dialog.is_delete == False
    )
    
    # 如果指定了对话类型，则添加类型过滤
    if dialog_type is not None:
        query = query.filter(Dialog.dialog_type == dialog_type)
    
    dialog: Dialog = query.first()
    return dialog


@with_session
def db_save_dialog(session, name, user_id, dialog_type=1, kb_ids=[]) -> Dialog:
    dialog = Dialog(name=name, user_id=user_id, dialog_type=dialog_type, kb_ids=kb_ids)
    session.add(dialog)
    session.commit()
    return dialog


@with_session
def db_delete_dialog_by_id(session, user_id, dialog_id):
    d = session.query(Dialog).filter(Dialog.user_id == user_id, Dialog.id == dialog_id, Dialog.is_delete.is_(False)).first()
    if d is None:
        return None

    d.is_delete = True
    session.add(d)
    session.commit()
    return True


@with_session
def db_update_dialog_by_id(session, user_id, dialog_id, name=None, kb_ids=None):
    d = session.query(Dialog).filter(Dialog.user_id == user_id, Dialog.id == dialog_id, Dialog.is_delete.is_(False)).first()
    if d is None:
        return None

    if name:
        d.name = name
    if kb_ids:
        if isinstance(kb_ids[0], uuid.UUID):
            d.kb_ids = [str(i) for i in kb_ids]
        else:
            d.kb_ids = kb_ids

    session.add(d)
    session.commit()
    return True


@with_session
def db_get_dialog_by_role(session, user, offset, page_size, dialog_type):
    if user.check_superman_role():
        query = session.query(Dialog).filter(
            Dialog.dialog_type == dialog_type,
            Dialog.is_delete == False
        )
    elif user.check_admin_role():
        # 获取当前用户所在的所有组ID
        subquery_groups = session.query(UserGroup.group).filter(UserGroup.user == user.id)
        # 获取这些组中的所有用户ID
        subquery_users = session.query(UserGroup.user).filter(UserGroup.group.in_(subquery_groups))
        # 查询符合条件的Dialog记录
        query = session.query(Dialog).filter(
            Dialog.user_id.in_(subquery_users),
            Dialog.is_delete == False,
            Dialog.dialog_type == dialog_type
        )
    else:
        query = session.query(Dialog).filter(
            Dialog.user_id == user.id,
            Dialog.dialog_type == dialog_type,
            Dialog.is_delete == False
        )

    dialogs = query.order_by(desc(Dialog.updated_time)) \
        .offset(offset) \
        .limit(page_size) \
        .all()

    total = query.count()
    return total, dialogs
